#pragma once
#include <map>
#include <utility>// for std::pair

template<typename K, typename V>
class CMap : public std::map<K,V> {
public:
	CMap() = default;

	/**
	* @warning do not misuse `insert` and `insert_or_assign`, the difference is:
	*          `insert(K,V)` if key alreay exists, nothing happens!
	*          `insert_or_assign(K,V)` C++17 function: if key already exists, value will be replaced! If c++17 not supported, use `operator[]` instead:`(*this)[key] = value;`
	*/
	void insert(const K& key, const V& value) {
		std::map<K,V>::insert_or_assign(key, value);// <=> (*this)[key] = value;
	}

	void remove(const K& key) {
		std::map<K,V>::erase(key);
	}

	V take(const K& key) {
		if (!contains(key)) {
			if constexpr (std::is_pointer_v<V>)
				return nullptr;
			V v;
			return v;
		}
		auto value = std::map<K,V>::at(key);
		std::map<K,V>::erase(key);
		return value;
	}

	bool contains(const K& key) const {
		return std::map<K, V>::find(key) != std::map<K, V>::end();
	}

	/**@note
	V& map[key]	: If the key exists, return a reference to the corresponding value;
				  If it does not exist, insert {key, V()} and return a reference to the new value
	V& map.at(key) : If the key exists, return a reference to the corresponding value;
					 If it does not exist, throw std:: out_of_range exception
	*/

	V value(const K& key) {
		if (!contains(key)) {
			if constexpr (std::is_pointer_v<V>)
				return nullptr;
			V v;
			return v;
		}
		return std::map<K,V>::at(key);
	}

	/**@note if key not exists in map, throw a std::out_of_range exception*/
	V& valueRef(const K& key) { return this->at(key); }

	std::vector<V> values() const {
		std::vector<V> result;
		for (const auto& pair : *this)
			result.push_back(pair.value);
		return result;
	}

	std::vector<K> keys() const {
		std::vector<K> result;
		for (const auto& pair : *this)
			result.push_back(pair.key);
		return result;
	}

	/**
	*@brief proxy the default Iterator: let Map's iterators return {key, value} directly instead of `std::pair`
	***********************************************************************************************************
	*/
	struct Proxy {
		const K& key;
		V& value;

		Proxy(std::pair<const K, V>& p) : key(p.first), value(p.second) {}
	};

	struct ConstProxy {
		const K& key;
		const V& value;

		ConstProxy(const std::pair<const K, V>& p) : key(p.first), value(p.second) {}
	};

	struct Iterator {
		typename std::map<K, V>::iterator it;

		Iterator(typename std::map<K, V>::iterator it) : it(it) {}

		Proxy operator*() {
			return Proxy(*it);
		}

		Iterator& operator++() {
			++it;
			return *this;
		}

		bool operator!=(const Iterator& other) const {
			return it != other.it;
		}
	};

	struct ConstIterator {
		typename std::map<K, V>::const_iterator it;

		ConstIterator(typename std::map<K, V>::const_iterator it) : it(it) {}

		ConstProxy operator*() const {
			return ConstProxy(*it);
		}

		ConstIterator& operator++() {
			++it;
			return *this;
		}

		bool operator!=(const ConstIterator& other) const {
			return it != other.it;
		}
	};

	// Non const versions of begin/end
	Iterator begin() { return Iterator(std::map<K, V>::begin()); }
	Iterator end() { return Iterator(std::map<K, V>::end()); }

	// Const version of begin/end, return const_iterator (for const objects)
	ConstIterator begin() const { return ConstIterator(std::map<K, V>::begin()); }
	ConstIterator end() const { return ConstIterator(std::map<K, V>::end()); }
};